package com.tbit.uqbike.client.filter.servlet;

import javax.servlet.ServletOutputStream;
import javax.servlet.WriteListener;
import javax.servlet.http.HttpServletResponse;
import javax.servlet.http.HttpServletResponseWrapper;
import java.io.ByteArrayOutputStream;
import java.io.IOException;
import java.io.OutputStreamWriter;
import java.io.PrintWriter;
import java.nio.charset.StandardCharsets;

public class CachedBodyHttpServletResponse extends HttpServletResponseWrapper {
    private final ByteArrayOutputStream outputStream = new ByteArrayOutputStream();
    private final ServletOutputStream servletOutputStream;
    private final PrintWriter printWriter;

    public CachedBodyHttpServletResponse(HttpServletResponse response) {
        super(response);
        this.servletOutputStream = new ServletOutputStream() {
            @Override
            public void write(int b) {
                outputStream.write(b);
            }

            @Override
            public boolean isReady() { return true; }

            @Override
            public void setWriteListener(WriteListener listener) {}
        };
        this.printWriter = new PrintWriter(new OutputStreamWriter(outputStream, StandardCharsets.UTF_8));
    }

    @Override
    public ServletOutputStream getOutputStream() {
        return servletOutputStream;
    }

    @Override
    public PrintWriter getWriter() {
        return printWriter;
    }

    public String getContent() {
        return new String(outputStream.toByteArray(), StandardCharsets.UTF_8);
    }

    public void copyBodyToResponse() throws IOException {
        PrintWriter writer = super.getWriter();
        writer.write(getContent());
        writer.flush();
    }
}